#include "SortPrimitive.h"
//#include "QD3DTimeQuery.h"
#include <sstream>

#ifdef QUERY_TIME
//#define SORT_QUERY_TIME
#endif


SortPrimitive::SortPrimitive (glm::uvec2 bufSize,
		QD3DResourcePack *pDataPack, UINT structSize,
		const std::string &structBufType) :
		_rp_dataPack (pDataPack), _structBufType (structBufType),
		_bufSize (bufSize.x * bufSize.y), _groupSize (bufSize.x),
		_groupNum (bufSize.y)
{
	QD3DUtils d3DUtils;

	_p_dataSizeDxConstBuf = d3DUtils.createConstBuffer< SortConstVars > ();

	ID3D11Buffer *pTempDxBuf = d3DUtils.createStructuredBuffer (
			_bufSize, structSize, NULL,
			D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_UNORDERED_ACCESS);
	_p_tempPack = new QD3DResourcePack (pTempDxBuf);

	QD3DShaderDesc shaderDesc;
	
	shaderDesc.setDefine ("TYPE", _structBufType);
	std::ostringstream sOut;
	sOut << _groupSize;
	shaderDesc.setDefine ("BITONIC_BLOCK_SIZE", sOut.str ());
	sOut = std::ostringstream ();
	sOut << TRANSPOSE_BLOCK_SIZE;
	shaderDesc.setDefine ("TRANSPOSE_BLOCK_SIZE", sOut.str ());
	_p_sortCS = d3DUtils.createShaderFromFile< ID3D11ComputeShader > (
			QD3DUtils::COMPUTE_SHADER, "cs_5_0",
			"..\\Data\\Shaders\\Primitives\\BitonicSortCS.hlsl",
			"BitonicSort", shaderDesc);
	_p_transposeCS = d3DUtils.createShaderFromFile< ID3D11ComputeShader > (
			QD3DUtils::COMPUTE_SHADER, "cs_5_0",
			"..\\Data\\Shaders\\Primitives\\BitonicSortCS.hlsl",
			"MatrixTranspose", shaderDesc);
}

SortPrimitive::~SortPrimitive ()
{
	SAFE_RELEASE (_p_dataSizeDxConstBuf);
	SAFE_RELEASE (_p_sortCS);
	SAFE_RELEASE (_p_transposeCS);
	delete _p_tempPack;
}

void SortPrimitive::run ()
{
	QD3DUtils d3DUtils;

#ifdef SORT_QUERY_TIME
	QD3DTimeQuery timeQuery;
	timeQuery.begin ();
	timeQuery.addCheckPoint ("start");
#endif

	SortConstVars constVars;

	QVector< ID3D11Buffer * > sortDxConstBufPtrs (1);
	QVector< ID3D11ShaderResourceView * > sortDxSRVPtrs (1);
	QVector< ID3D11UnorderedAccessView * > sortDxUAVPtrs (1);
	sortDxConstBufPtrs[0] = _p_dataSizeDxConstBuf;
	d3DUtils.setConstBufs (QD3DUtils::COMPUTE_SHADER, 0, sortDxConstBufPtrs);

	// Sort the data
	d3DUtils.setComputeShader (_p_sortCS);
	sortDxUAVPtrs[0] = _rp_dataPack->getUAV ();
	d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);

	// First sort the rows for the levels <= to the block size
	for( UINT level = 2 ; level <= _groupSize ; level = level * 2 ) {
		constVars.var.setData (level, level, _groupNum, _groupSize);
		d3DUtils.updateSubresource (_p_dataSizeDxConstBuf, 0, &constVars);
		d3DUtils.dispatch (_bufSize / _groupSize, 1, 1);
	}
	d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());

	// Then sort the rows and columns for the levels > than the block size
	// Transpose. Sort the Columns. Transpose. Sort the Rows.
#ifdef SORT_QUERY_TIME
	timeQuery.addCheckPoint ("step 2");
#endif // QUERY_TIME
	for (UINT level = (_groupSize * 2); level <= _bufSize;
			level = level * 2) {
		constVars.var.setData ((level / _groupSize),
				(level & ~_bufSize) / _groupSize,
				_groupSize, _groupNum);
		d3DUtils.updateSubresource (_p_dataSizeDxConstBuf, 0, &constVars);

		// Transpose the data from buffer 1 into buffer 2
		sortDxSRVPtrs[0] = _rp_dataPack->getSRV ();
		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
		sortDxUAVPtrs[0] = _p_tempPack->getUAV ();
		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
		d3DUtils.setComputeShader (_p_transposeCS);
		d3DUtils.dispatch (_groupSize / TRANSPOSE_BLOCK_SIZE,
				_groupNum / TRANSPOSE_BLOCK_SIZE, 1);
		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0,
				sortDxSRVPtrs.size ());

		// Sort the transposed column data
		d3DUtils.setComputeShader (_p_sortCS);
		d3DUtils.dispatch (_bufSize / _groupSize, 1, 1);
		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0,
				sortDxUAVPtrs.size ());

		constVars.var.setData (_groupSize, level,
				_groupNum, _groupSize);
		d3DUtils.updateSubresource (_p_dataSizeDxConstBuf, 0, &constVars);

		// Transpose the data from buffer 2 back into buffer 1
		sortDxSRVPtrs[0] = _p_tempPack->getSRV ();
		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
		sortDxUAVPtrs[0] = _rp_dataPack->getUAV ();
		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
		d3DUtils.setComputeShader (_p_transposeCS);
		d3DUtils.dispatch (_groupNum / TRANSPOSE_BLOCK_SIZE,
				_groupSize / TRANSPOSE_BLOCK_SIZE, 1);
		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0,
				sortDxSRVPtrs.size ());

		// Sort the row data
		d3DUtils.setComputeShader (_p_sortCS);
		d3DUtils.dispatch (_bufSize / _groupSize, 1, 1);
		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0,
				sortDxUAVPtrs.size ());
	}

// 	d3DUtils.setComputeShader (_pSortCS);
// 	
// 	constVars.var.groupSize = groupSizeArr[0][1];
// 	constVars.var.groupThreadSize = groupSizeArr[0][2];
// 	constVars.var.minLevelSize = 2;
// 	constVars.var.maxLevelSize = constVars.var.groupThreadSize * 2;
// 	constVars.var.useExtLevelMask = 0;
// 	d3DUtils.updateSubresource (_pDataSizeDxConstBuf, 0, &constVars);
// 
// 	sortDxSRVPtrs[0] = NULL;
// 	d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
// 	sortDxUAVPtrs[0] = _pDataPack->getUAV ();
// 	d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
// 
// 	timeQuery.addCheckPoint ("start first");
// 	d3DUtils.dispatch (groupSizeArr[0][0] / groupSizeArr[0][1], 1, 1);
// 	timeQuery.addCheckPoint ("end first");
// 
// 	d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs.size ());
// 	d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());
// 
// 	for (UINT levelSize = groupSizeArr[0][1] * 2;
// 			levelSize <= groupSizeArr[0][0]; levelSize <<= 1) {
// 		TransposeConstVars transposeConstVars;
// 
// 		// Transpose
// 		d3DUtils.setComputeShader (_pTransposeCS);
// 		transposeConstVars.var.groupSize = groupSizeArr[0][1];
// 		transposeConstVars.var.groupNum = groupSizeArr[0][0] / groupSizeArr[0][1];
// 		d3DUtils.updateSubresource (_pDataSizeDxConstBuf, 0, &transposeConstVars);
// 
// 		sortDxSRVPtrs[0] = _pDataPack->getSRV ();
// 		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
// 		sortDxUAVPtrs[0] = _pTempPack->getUAV ();
// 		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
// 		d3DUtils.dispatch (transposeConstVars.var.groupSize / TRANSPOSE_BLOCK_SIZE,
// 				transposeConstVars.var.groupNum / TRANSPOSE_BLOCK_SIZE, 1);
// 		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs.size ());
// 		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());
// 
// 		// Sort column
// 		d3DUtils.setComputeShader (_pSortCS);
// 		constVars.var.groupSize = groupSizeArr[0][1];
// 		constVars.var.groupThreadSize = groupSizeArr[0][2];
// 		constVars.var.minLevelSize = levelSize / groupSizeArr[0][1];
// 		constVars.var.maxLevelSize = levelSize / groupSizeArr[0][1];
// 		constVars.var.extLevelMask = (levelSize == groupSizeArr[0][0] ? 0 : levelSize / groupSizeArr[0][1]);
// 		constVars.var.useExtLevelMask = 1;
// 		d3DUtils.updateSubresource (_pDataSizeDxConstBuf, 0, &constVars);
// 
// 		sortDxSRVPtrs[0] = NULL;
// 		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
// 		sortDxUAVPtrs[0] = _pTempPack->getUAV ();
// 		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
// 
// 		d3DUtils.dispatch (groupSizeArr[0][0] / groupSizeArr[0][1], 1, 1);
// 
// 		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs.size ());
// 		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());
// 
// 		// Transpose back
// 		d3DUtils.setComputeShader (_pTransposeCS);
// 		transposeConstVars.var.groupSize = groupSizeArr[0][0] / groupSizeArr[0][1];
// 		transposeConstVars.var.groupNum = groupSizeArr[0][1];
// 		d3DUtils.updateSubresource (_pDataSizeDxConstBuf, 0, &transposeConstVars);
// 
// 		sortDxSRVPtrs[0] = _pTempPack->getSRV ();
// 		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
// 		sortDxUAVPtrs[0] = _pDataPack->getUAV ();
// 		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
// 		d3DUtils.dispatch (transposeConstVars.var.groupSize / TRANSPOSE_BLOCK_SIZE,
// 				transposeConstVars.var.groupNum / TRANSPOSE_BLOCK_SIZE, 1);
// 		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs.size ());
// 		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());
// 
// 		// Sort row
// 		d3DUtils.setComputeShader (_pSortCS);
// 		constVars.var.groupSize = groupSizeArr[0][1];
// 		constVars.var.groupThreadSize = groupSizeArr[0][2];
// 		constVars.var.minLevelSize = constVars.var.groupThreadSize * 2;
// 		constVars.var.maxLevelSize = constVars.var.groupThreadSize * 2;
// 		constVars.var.extLevelMask = levelSize / groupSizeArr[0][1] * groupSizeArr[0][2] * 2;
// 		constVars.var.useExtLevelMask = 1;
// 		d3DUtils.updateSubresource (_pDataSizeDxConstBuf, 0, &constVars);
// 
// 		sortDxSRVPtrs[0] = NULL;
// 		d3DUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs);
// 		sortDxUAVPtrs[0] = _pDataPack->getUAV ();
// 		d3DUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs);
// 
// 		d3DUtils.dispatch (groupSizeArr[0][0] / groupSizeArr[0][1], 1, 1);
// 
// 		d3DUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxSRVPtrs.size ());
// 		d3DUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, sortDxUAVPtrs.size ());
// 	}
	d3DUtils.clearConstBufs (QD3DUtils::COMPUTE_SHADER, 0,
			sortDxConstBufPtrs.size ());
#ifdef SORT_QUERY_TIME
	timeQuery.addCheckPoint ("end");
	timeQuery.end ();

	std::ostringstream sOut;
	sOut << "Sort: " << timeQuery.getDuration ("start", "end") << "ms\n";
 	sOut << "Step2: " << timeQuery.getDuration ("step 2", "end") << "ms\n";
	OutputDebugStringA (sOut.str ().c_str ());
#endif
}